import sys
import os
import re
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen3-4B"

# load the tokenizer and the model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)

def convert_fortran_to_c(fortran_code):
    # Prepare the prompt for Fortran to C conversion
    prompt = f"Convert the following Fortran code to C code. Ensure proper handling of array declarations, loops, and function calls. Use 0-based indexing for arrays in C. Here is the Fortran code:\n\n{fortran_code}\n\nProvide the converted C code and explain your conversion process."
    messages = [
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # Conduct text completion
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=32768
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

    # Parsing thinking content
    try:
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0

    thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

    print("Thinking process:", thinking_content)
    print("Converted C code:", content)
    return content

def main():
    if len(sys.argv) < 2:
        print("Usage: python autoF2C-qwen.py <input_file.f90>")
        sys.exit(1)

    input_file = sys.argv[1]
    if not os.path.exists(input_file):
        print(f"Error: Input file {input_file} not found.")
        sys.exit(1)

    # Read the Fortran file
    with open(input_file, 'r') as f:
        fortran_code = f.read()

    # Convert using Qwen model
    output_content = convert_fortran_to_c(fortran_code)

    # Separate C code and markdown content
    c_code = ""
    markdown_content = ""
    in_c_code_block = False
    c_code_lines = []
    markdown_lines = []

    for line in output_content.splitlines():
        if line.strip().startswith('```c'):
            in_c_code_block = True
            continue
        elif line.strip().startswith('```') and in_c_code_block:
            in_c_code_block = False
            continue
        
        if in_c_code_block:
            c_code_lines.append(line)
        else:
            markdown_lines.append(line)

    c_code = "\n".join(c_code_lines).strip()
    markdown_content = "\n".join(markdown_lines).strip()

    # Write the output to a .c file
    output_file_c = os.path.splitext(input_file)[0] + '.c'
    with open(output_file_c, 'w') as f:
        f.write(c_code)

    # Write the markdown explanations to a .md file
    output_file_md = os.path.splitext(input_file)[0] + '_explanation.md'
    with open(output_file_md, 'w') as f:
        f.write(markdown_content)

    print(f"Conversion complete. C code written to {output_file_c}")
    print(f"Explanations written to {output_file_md}")

if __name__ == "__main__":
    main()
